"""
配置管理模块
"""
import json
import os
try:
    import yaml
    YAML_AVAILABLE = True
except ImportError:
    YAML_AVAILABLE = False
    print("警告: PyYAML 未安装，无法加载YAML配置文件。请安装: pip install PyYAML")


class Config:
    """训练配置类"""
    
    def __init__(self, config_dict=None):
        """
        初始化配置
        
        Args:
            config_dict (dict): 配置字典
        """
        # 默认配置
        self.default_config = {
            # 数据相关
            'data_dir': '../dataset/Mini-ImageNet-Dataset',
            'label_file': '../dataset/label.json',
            'batch_size': 256,
            'num_workers': 0,
            'pin_memory': True,
            'img_size': 224,
            'num_classes': 100,
            
            # 训练相关
            'num_epochs': 30,
            'lr': 0.01,
            'momentum': 0.9,
            'weight_decay': 1e-4,
            'print_freq': 50,
            
            # 学习率调度器
            'scheduler_type': 'StepLR',
            'step_size': 10,
            'gamma': 0.1,
            
            # 模型保存和日志
            'log_dir': '../logs',
            'model_save_dir': '../models',
            'save_best': True,
            
            # 数据增强
            'normalize_mean': [0.485, 0.456, 0.406],
            'normalize_std': [0.229, 0.224, 0.225],
            'random_horizontal_flip': 0.5,
            'random_rotation': 10,
        }
        
        # 更新配置
        if config_dict:
            self.default_config.update(config_dict)
    
    def get(self, key, default=None):
        """获取配置值"""
        return self.default_config.get(key, default)
    
    def set(self, key, value):
        """设置配置值"""
        self.default_config[key] = value
    
    def update(self, config_dict):
        """更新配置"""
        self.default_config.update(config_dict)
    
    def save_config(self, file_path):
        """保存配置到文件（支持JSON和YAML格式）"""
        os.makedirs(os.path.dirname(file_path), exist_ok=True)
        
        if file_path.endswith('.yaml') or file_path.endswith('.yml'):
            if YAML_AVAILABLE:
                with open(file_path, 'w', encoding='utf-8') as f:
                    yaml.dump(self.default_config, f, default_flow_style=False, 
                            allow_unicode=True, indent=2)
            else:
                print("PyYAML 未安装，将保存为JSON格式")
                json_path = file_path.replace('.yaml', '.json').replace('.yml', '.json')
                with open(json_path, 'w', encoding='utf-8') as f:
                    json.dump(self.default_config, f, indent=4, ensure_ascii=False)
        else:
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(self.default_config, f, indent=4, ensure_ascii=False)
    
    def load_config(self, file_path):
        """从文件加载配置（支持JSON和YAML格式）"""
        if os.path.exists(file_path):
            if file_path.endswith('.yaml') or file_path.endswith('.yml'):
                if YAML_AVAILABLE:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        config_dict = yaml.safe_load(f)
                        self.default_config.update(self._flatten_dict(config_dict))
                else:
                    print("PyYAML 未安装，无法加载YAML文件")
            else:
                with open(file_path, 'r', encoding='utf-8') as f:
                    config_dict = json.load(f)
                    self.default_config.update(config_dict)
        else:
            print(f"配置文件 {file_path} 不存在，使用默认配置")
    
    def _flatten_dict(self, nested_dict, parent_key='', sep='_'):
        """将嵌套字典展平"""
        items = []
        for k, v in nested_dict.items():
            new_key = f"{parent_key}{sep}{k}" if parent_key else k
            if isinstance(v, dict):
                items.extend(self._flatten_dict(v, new_key, sep=sep).items())
            else:
                items.append((new_key, v))
        return dict(items)
    
    def load_yaml_config(self, file_path):
        """加载YAML配置文件（保持嵌套结构）"""
        if not YAML_AVAILABLE:
            print("PyYAML 未安装，无法加载YAML文件")
            return {}
        
        if os.path.exists(file_path):
            with open(file_path, 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        else:
            print(f"配置文件 {file_path} 不存在")
            return {}
    
    def get_model_save_path(self, model_name):
        """获取模型保存路径"""
        model_dir = os.path.join(self.get('model_save_dir'), model_name)
        os.makedirs(model_dir, exist_ok=True)
        return os.path.join(model_dir, f'{model_name}_model.pth')
    
    def get_log_dir(self, model_name):
        """获取日志目录"""
        return os.path.join(self.get('log_dir'), model_name)
    
    def print_config(self):
        """打印配置信息"""
        print("=== 当前配置 ===")
        for key, value in self.default_config.items():
            print(f"{key}: {value}")
        print("=" * 20)
    
    def to_dict(self):
        """转换为字典"""
        return self.default_config.copy()


# 预定义配置
class AlexNetConfig(Config):
    """AlexNet专用配置"""
    
    def __init__(self):
        alexnet_config = {
            'lr': 0.01,
            'momentum': 0.9,
            'weight_decay': 1e-4,
            'scheduler_type': 'StepLR',
            'step_size': 10,
            'gamma': 0.1,
            'img_size': 224,
            'batch_size': 256,
            'num_epochs': 30,
        }
        super().__init__(alexnet_config)


class VGGConfig(Config):
    """VGG专用配置"""
    
    def __init__(self):
        vgg_config = {
            'lr': 0.001,
            'momentum': 0.9,
            'weight_decay': 1e-4,
            'scheduler_type': 'StepLR',
            'step_size': 7,
            'gamma': 0.1,
            'img_size': 224,
            'batch_size': 128,
            'num_epochs': 50,
        }
        super().__init__(vgg_config)


class ResNetConfig(Config):
    """ResNet专用配置"""
    
    def __init__(self):
        resnet_config = {
            'lr': 0.1,
            'momentum': 0.9,
            'weight_decay': 1e-4,
            'scheduler_type': 'StepLR',
            'step_size': 30,
            'gamma': 0.1,
            'img_size': 224,
            'batch_size': 256,
            'num_epochs': 100,
        }
        super().__init__(resnet_config)


class YAMLConfig(Config):
    """YAML配置文件加载器"""
    
    def __init__(self, yaml_file=None):
        """
        从YAML文件初始化配置
        
        Args:
            yaml_file (str): YAML配置文件路径
        """
        super().__init__()
        
        if yaml_file:
            self.load_from_yaml(yaml_file)
    
    def load_from_yaml(self, yaml_file):
        """从YAML文件加载配置"""
        if not YAML_AVAILABLE:
            print("PyYAML 未安装，无法加载YAML文件")
            return
        
        yaml_config = self.load_yaml_config(yaml_file)
        if yaml_config:
            # 将嵌套的YAML配置转换为扁平化的配置
            flattened_config = self._convert_yaml_to_flat(yaml_config)
            self.default_config.update(flattened_config)
            print(f"已从 {yaml_file} 加载配置")
    
    def _convert_yaml_to_flat(self, yaml_config):
        """将YAML的嵌套结构转换为扁平化配置"""
        flat_config = {}
        
        # 数据相关配置
        if 'data' in yaml_config:
            data_config = yaml_config['data']
            flat_config.update({
                'data_dir': data_config.get('data_dir', '../dataset/Mini-ImageNet-Dataset'),
                'label_file': data_config.get('label_file', '../dataset/label.json'),
                'batch_size': data_config.get('batch_size', 256),
                'num_workers': data_config.get('num_workers', 0),
                'pin_memory': data_config.get('pin_memory', True),
                'img_size': data_config.get('img_size', 224),
                'num_classes': data_config.get('num_classes', 100),
            })
        
        # 数据增强配置
        if 'augmentation' in yaml_config:
            aug_config = yaml_config['augmentation']
            flat_config.update({
                'normalize_mean': aug_config.get('normalize_mean', [0.485, 0.456, 0.406]),
                'normalize_std': aug_config.get('normalize_std', [0.229, 0.224, 0.225]),
                'random_horizontal_flip': aug_config.get('random_horizontal_flip', 0.5),
                'random_rotation': aug_config.get('random_rotation', 10),
            })
        
        # 训练配置
        if 'training' in yaml_config:
            train_config = yaml_config['training']
            flat_config.update({
                'num_epochs': train_config.get('num_epochs', 30),
                'lr': train_config.get('lr', 0.01),
                'momentum': train_config.get('momentum', 0.9),
                'weight_decay': train_config.get('weight_decay', 1e-4),
                'print_freq': train_config.get('print_freq', 50),
            })
        
        # 学习率调度器配置
        if 'scheduler' in yaml_config:
            sched_config = yaml_config['scheduler']
            flat_config.update({
                'scheduler_type': sched_config.get('type', 'StepLR'),
                'step_size': sched_config.get('step_size', 10),
                'gamma': sched_config.get('gamma', 0.1),
                'T_max': sched_config.get('T_max', 200),  # 用于CosineAnnealingLR
            })
        
        # 路径配置
        if 'paths' in yaml_config:
            path_config = yaml_config['paths']
            flat_config.update({
                'log_dir': path_config.get('log_dir', '../logs'),
                'model_save_path': path_config.get('model_save_path', '../models/model.pth'),
                'save_best': path_config.get('save_best', True),
            })
        
        # 模型配置
        if 'model' in yaml_config:
            model_config = yaml_config['model']
            flat_config.update({
                'model_name': model_config.get('name', 'Model'),
                'model_type': model_config.get('model_type', 'resnet50'),
                'dropout': model_config.get('dropout', 0.5),
                'use_warmup': model_config.get('use_warmup', True),
            })
        
        return flat_config